#include "stdafx.h"
#include "net/WebSocket.hpp"
#include "net/WebSocketImpl.hpp"
#include "net/HTTPRequest.hpp"
#include "net/HTTPResponse.hpp"
#include "net/HTTPServerRequest.hpp"
#include "net/HTTPServerRequestImpl.hpp"
#include "net/HTTPServerResponse.hpp"
#include "net/HTTPServerResponseImpl.hpp"
#include "net/HTTPClientSession.hpp"
#include "stream/streambuf.hpp"
#include "stream/buffer.hpp"
#include "stream/StreamCopier.hpp"
#include "stream/base64encode.hpp"
using namespace JHCPP::stream;
#include "code/sha1engine.hpp"
#include "code/digestengine.hpp"
using namespace JHCPP::code;

#include <sstream>


NET_NAMESPACE_BEGIN

const std::string CWebSocket::WEBSOCKET_GUID("258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
const std::string CWebSocket::WEBSOCKET_VERSION("13");
CHTTPCredentials CWebSocket::m_defaultCreds;

CWebSocket::CWebSocket(CHTTPServerRequest& request, CHTTPServerResponse& response) : CStreamSocket(accept(request, response))
{
}

CWebSocket::CWebSocket(CHTTPClientSession& cs, CHTTPRequest& request, CHTTPResponse& response)
	: CStreamSocket(connect(cs, request, response, m_defaultCreds))
{
}

CWebSocket::CWebSocket(CHTTPClientSession& cs, CHTTPRequest& request, CHTTPResponse& response, CHTTPCredentials& credentials)
	: CStreamSocket(connect(cs, request, response, credentials))
{
}

CWebSocket::CWebSocket(const CSocket& socket) : CStreamSocket(socket)
{
	if (!dynamic_cast<CWebSocketImpl*>(impl()))
		throw InvalidArgumentException("Cannot assign incompatible socket");
}

CWebSocket::~CWebSocket()
{
}

CWebSocket& CWebSocket::operator = (const CSocket& socket)
{
	if (dynamic_cast<CWebSocketImpl*>(socket.impl()))
		CSocket::operator = (socket);
	else
		throw InvalidArgumentException("Cannot assign incompatible socket");
	return *this;
}

void CWebSocket::shutdown()
{
	shutdown(WS_NORMAL_CLOSE);
}

void CWebSocket::shutdown(UInt16 statusCode, const std::string& statusMessage)
{
	TBuffer<char> buffer(statusMessage.size() + 2);
	CMemoryOutputStream ostr(buffer.begin(), buffer.size());
	CBinaryWriter writer(ostr, CBinaryWriter::NETWORK_BYTE_ORDER);
	writer << statusCode;
	writer.writeRaw(statusMessage);
	sendFrame(buffer.begin(), static_cast<int>(ostr.charsWritten()), FRAME_FLAG_FIN | FRAME_OP_CLOSE);
}

int CWebSocket::sendFrame(const void* buffer, int length, int flags)
{
	flags |= FRAME_OP_SETRAW;
	return static_cast<CWebSocketImpl*>(impl())->sendBytes(buffer, length, flags);
}

int CWebSocket::receiveFrame(void* buffer, int length, int& flags)
{
	int n = static_cast<CWebSocketImpl*>(impl())->receiveBytes(buffer, length, 0);
	flags = static_cast<CWebSocketImpl*>(impl())->frameFlags();
	return n;
}

CWebSocket::Mode CWebSocket::mode() const
{
	return static_cast<CWebSocketImpl*>(impl())->mustMaskPayload() ? WS_CLIENT : WS_SERVER;
}

CWebSocketImpl* CWebSocket::accept(CHTTPServerRequest& request, CHTTPServerResponse& response)
{
	if (request.hasToken("Connection", "upgrade") && icompare(request.get("Upgrade", ""), "websocket") == 0)
	{
		std::string version = request.get("Sec-WebSocket-Version", "");
		if (version.empty()) throw WebSocketException("Missing Sec-WebSocket-Version in handshake request", WS_ERR_HANDSHAKE_NO_VERSION);
		if (version != WEBSOCKET_VERSION) throw WebSocketException("Unsupported WebSocket version requested", version, WS_ERR_HANDSHAKE_UNSUPPORTED_VERSION);
		std::string key = request.get("Sec-WebSocket-Key", "");
		trimInPlace(key);
		if (key.empty()) throw WebSocketException("Missing Sec-WebSocket-Key in handshake request", WS_ERR_HANDSHAKE_NO_KEY);

		response.setStatusAndReason(CHTTPResponse::HTTP_SWITCHING_PROTOCOLS);
		response.set("Upgrade", "websocket");
		response.set("Connection", "Upgrade");
		response.set("Sec-WebSocket-Accept", computeAccept(key));
		response.setContentLength(0);
		response.send().flush();
		return new CWebSocketImpl(static_cast<CStreamSocketImpl*>(static_cast<CHTTPServerRequestImpl&>(request).detachSocket().impl()), false);
	}
	else throw WebSocketException("No WebSocket handshake", WS_ERR_NO_HANDSHAKE);
}

CWebSocketImpl* CWebSocket::connect(CHTTPClientSession& cs, CHTTPRequest& request, CHTTPResponse& response, CHTTPCredentials& credentials)
{
	if (!cs.getProxyHost().empty() && !cs.secure())
	{
		cs.proxyTunnel();
	}
	std::string key = createKey();
	request.set("Connection", "Upgrade");
	request.set("Upgrade", "websocket");
	request.set("Sec-WebSocket-Version", WEBSOCKET_VERSION);
	request.set("Sec-WebSocket-Key", key);
	request.setChunkedTransferEncoding(false);
	cs.setKeepAlive(true);
	cs.sendRequest(request);
	std::istream& istr = cs.receiveResponse(response);
	if (response.getStatus() == CHTTPResponse::HTTP_SWITCHING_PROTOCOLS)
	{
		return completeHandshake(cs, response, key);
	}
	else if (response.getStatus() == CHTTPResponse::HTTP_UNAUTHORIZED)
	{
		CNullOutputStream null;
		CStreamCopier::copyStream(istr, null);
		credentials.authenticate(request, response);
		if (!cs.getProxyHost().empty() && !cs.secure())
		{
			cs.reset();
			cs.proxyTunnel();
		}
		cs.sendRequest(request);
		cs.receiveResponse(response);
		if (response.getStatus() == CHTTPResponse::HTTP_SWITCHING_PROTOCOLS)
		{
			return completeHandshake(cs, response, key);
		}
		else if (response.getStatus() == CHTTPResponse::HTTP_UNAUTHORIZED)
		{
			throw WebSocketException("Not authorized", WS_ERR_UNAUTHORIZED);
		}
	}
	if (response.getStatus() == CHTTPResponse::HTTP_OK)
	{
		throw WebSocketException("The server does not understand the WebSocket protocol", WS_ERR_NO_HANDSHAKE);
	}
	else
	{
		throw WebSocketException("Cannot upgrade to WebSocket connection", response.getReason(), WS_ERR_NO_HANDSHAKE);
	}
}

CWebSocketImpl* CWebSocket::completeHandshake(CHTTPClientSession& cs, CHTTPResponse& response, const std::string& key)
{
	std::string connection = response.get("Connection", "");
	if (icompare(connection, "Upgrade") != 0) 
		throw WebSocketException("No Connection: Upgrade header in handshake response", WS_ERR_NO_HANDSHAKE);
	std::string upgrade = response.get("Upgrade", "");
	if (icompare(upgrade, "websocket") != 0)
		throw WebSocketException("No Upgrade: websocket header in handshake response", WS_ERR_NO_HANDSHAKE);
	std::string accept = response.get("Sec-WebSocket-Accept", "");
	if (accept != computeAccept(key))
		throw WebSocketException("Invalid or missing Sec-WebSocket-Accept header in handshake response", WS_ERR_NO_HANDSHAKE);
	return new CWebSocketImpl(static_cast<CStreamSocketImpl*>(cs.detachSocket().impl()), true);
}

std::string CWebSocket::createKey()
{
	CRandom rnd;
	std::ostringstream ostr;
	CBase64Encode base64(ostr);
	CBinaryWriter writer(base64);
	writer << rnd.next() << rnd.next() << rnd.next() << rnd.next();
	base64.close();
	return ostr.str();
}

std::string CWebSocket::computeAccept(const std::string& key)
{
	std::string accept(key);
	accept += WEBSOCKET_GUID;
	CSHA1Engine sha1;
	sha1.update(accept);
	CDigestEngine::Digest d = sha1.digest();
	std::ostringstream ostr;
	CBase64Encode base64(ostr);
	base64.write(reinterpret_cast<const char*>(&d[0]), d.size());
	base64.close();
	return ostr.str();
}

NET_NAMESPACE_END
